#include "DXUT.h"
#define _USE_MATH_DEFINES
#include <cmath>
#include <stdlib.h>
#include <iostream>
#include "PDSampler/PDSampling.h"
#include "GenRandomSeq.h"

const float CGenRandomSeq::minrnd = 0.000001;
const float CGenRandomSeq::maxrnd = 1-minrnd;

CGenRandomSeq::CGenRandomSeq(void)
{
	m_currSequenceType = FOLDED_HAMMERSLEY;
	m_iNumSamples = 40;
	m_iNumLafortuneSmpls = 0;
	m_iNumWardSmpls = 0;

	ReleaseMemory();
	ReserveMemory();
}

CGenRandomSeq::CGenRandomSeq(SequenceType sType, int iNum)
{ 
	m_currSequenceType = sType;
	m_iNumSamples = iNum; 
	m_iNumLafortuneSmpls = 0;
	m_iNumWardSmpls = 0;

	ReleaseMemory();
	ReserveMemory();
}

// Inversion code courtesy of Matt Pharr and Greg Humphries
inline double CGenRandomSeq::FoldedRadicalInverse(int n, int base) {
	double val = 0;
	double invBase = 1.f/base, invBi = invBase;
	int modOffset = 0;
	while (val + base * invBi != val) {
		// Compute next digit of folded radical inverse
		int digit = ((n+modOffset) % base);
		val += digit * invBi;
		n /= base;
		invBi *= invBase;
		++modOffset;
	}
	return val;
}

inline double CGenRandomSeq::RadicalInverse(int n, int base) {
	double val = 0;
	double invBase = 1. / base, invBi = invBase;
	while (n > 0) {
		// Compute next digit of radical inverse
		int d_i = (n % base);
		val += d_i * invBi;
		n /= base;
		invBi *= invBase;
	}
	return val;
}

void CGenRandomSeq::genSequence() 
{
	switch (m_currSequenceType) {
	case HALTON:			genHaltonSequence();			break;
	case FOLDED_HALTON:		genFoldedHaltonSequence();		break;
	case HAMMERSLEY:		genHammersleySequence();		break;
	case FOLDED_HAMMERSLEY:	genFoldedHammersleySequence();	break;
	case POSSION_DISK:		genPossionDiskSequence();		break;
	case BEST_CANDIDATE:	genBestCandidateSequence();	break;
	case PENROSE:			genPenroseSequence();			break;
	}
}

void CGenRandomSeq::genHaltonSequence() 
{
	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = (float) RadicalInverse(i,2);
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		double e2 = 2.*M_PI*RadicalInverse(i,3);
		randcos[i] = (float) cos(e2);
		randsin[i] = (float) sin(e2);		
	}
}

void CGenRandomSeq::genFoldedHaltonSequence() 
{
	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = (float) FoldedRadicalInverse(i,2);
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		double e2 = 2.*M_PI*FoldedRadicalInverse(i,3);
		randcos[i] = (float) cos(e2);
		randsin[i] = (float) sin(e2);		
	}
}

void CGenRandomSeq::genHammersleySequence() 
{
	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = ((float) i)/((float) m_iNumSamples);
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		double e2 = 2.*M_PI*RadicalInverse(i,2);
		randcos[i] = (float) cos(e2);
		randsin[i] = (float) sin(e2);		
	}
}

void CGenRandomSeq::genFoldedHammersleySequence()
{
	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = ((float) i)/((float) m_iNumSamples);
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		double e2 = 2.*M_PI*FoldedRadicalInverse(i,2);
		randcos[i] = (float) cos(e2);
		randsin[i] = (float) sin(e2);		
	}
}

void CGenRandomSeq::genPossionDiskSequence() 
{
	float radius = 0.9f/sqrtf((float) m_iNumSamples);
	radius = min(radius, 0.2f);
	PDSampler *sampler = new PureSampler(radius);

	sampler->complete();

	while ((sampler->points.size()-m_iNumSamples) > 1000000)
		sampler->complete();

	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = sampler->points[i].x/2.f+0.5f;
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		float e2 = 2.f*((float) M_PI)*(sampler->points[i].y/2.f+0.5f);
		randcos[i] = (float) cosf(e2);
		randsin[i] = (float) sinf(e2);
	}

	SAFE_DELETE(sampler);
}

void CGenRandomSeq::genBestCandidateSequence() 
{
	float radius = 0.75f/sqrtf((float) m_iNumSamples);
	radius = min(radius, 0.2f);
	PDSampler *sampler = new BestCandidate(radius, true, 1);

	sampler->complete();

	while ((sampler->points.size()-m_iNumSamples) > 1000000)
		sampler->complete();

	for (int i=0; i < m_iNumSamples; i++) 
	{
		randnum[i] = sampler->points[i].x/2.f+0.5f;
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		float e2 = 2.f*((float) M_PI)*(sampler->points[i].y/2.f+0.5f);
		randcos[i] = cosf(e2);
		randsin[i] = sinf(e2);
	}

	SAFE_DELETE(sampler);
}

void CGenRandomSeq::genPenroseSequence() 
{
	float radius = 0.9f/sqrtf((float) m_iNumSamples);
	radius = min(radius, 0.2f);
	PDSampler *sampler = new BestCandidate(radius, true, 1);

	sampler->complete();

	while ((sampler->points.size()-m_iNumSamples) > 1000000)
		sampler->complete();

	for (int i=0; i < m_iNumSamples; i++) {
		randnum[i] = sampler->points[i].x/2.f+0.5f;
		if(randnum[i]<minrnd) randnum[i]=minrnd;
		if(randnum[i]>maxrnd) randnum[i]=maxrnd;
		randlog[i] = logf(randnum[i]);

		float e2 = 2.f*((float) M_PI)*(sampler->points[i].y/2.f+0.5f);
		randcos[i] = cosf(e2);
		randsin[i] = sinf(e2);
	}

	delete sampler;
}

// precompute the samples from the quasi-random values
void CGenRandomSeq::genLafortuneSamples(float cxy, float cz, float n) 
{
	if (m_iNumSamples > m_iNumLafortuneSmpls) 
	{
		LafortuneSmpls.RemoveAll();
		LafortuneScales.RemoveAll();
		m_iNumLafortuneSmpls = m_iNumSamples;
		for (unsigned int i = 0; i < m_iNumLafortuneSmpls; i++)
		{
			LafortuneSmpls.Add(D3DXVECTOR4(0,0,0,0));
			LafortuneScales.Add(0);
		}		
	}

	// format of the data
	// - u_x, u_y, u_z, biasCmp
	// - normalize*BRDF/PDF (partial)
	// Log[2, I*I/N/(2Pi)]*0.5 = Log[2, 512*512/40/(2*Pi)]*0.5 = 5.01329+1
	float lodPreComp = logf(512.f*512.f/((float) m_iNumSamples)/(2.f*((float) M_PI)))/((float) M_LN2)*0.5f+1.f;

	for (int i=0; i < m_iNumSamples; i++) {
		float costheta = powf(randnum[i], 1.f/(n+1.f));
		float sintheta = sqrtf(1.f - costheta*costheta);

		LafortuneSmpls[i].x = randcos[i]*sintheta;
		LafortuneSmpls[i].y = randsin[i]*sintheta;
		LafortuneSmpls[i].z = costheta;

		float pdf = (n+1)*powf(costheta, n)/(((float) M_PI)*2.f);

		LafortuneSmpls[i].w = max(0.f, lodPreComp - logf(pdf)/((float) M_LN2)*0.5f);
		LafortuneScales[i] = (n+2)*powf(costheta, n)/pdf;
	}
}

void CGenRandomSeq::genWardSamples(float alphax, float alphay) 
{
	if (m_iNumSamples > m_iNumWardSmpls) 
	{
		WardSmpls.RemoveAll();
		WardScales.RemoveAll();
		m_iNumWardSmpls = m_iNumSamples;
		for (unsigned int i = 0; i < m_iNumWardSmpls; i++)
		{
			WardSmpls.Add(D3DXVECTOR4(0,0,0,0));
			WardScales.Add(0);
		}		
		
	}

	for (int i=0; i < m_iNumSamples; i++) {
		float cosp = randcos[i] * alphax;
		float sinp = randsin[i] * alphay;

		float d = 1.f/sqrtf(cosp*cosp + sinp*sinp);
		cosp *= d;
		sinp *= d;

		d = -randlog[i] /( cosp*cosp/alphax/alphax + sinp*sinp/alphay/alphay);
		float hz = sqrtf(1.f/(d+1.f));
		float sint = sqrt(d)*hz;

		WardSmpls[i].x = sint*cosp;
		WardSmpls[i].y = sint*sinp;
		WardSmpls[i].z = hz;

		// compute the partial PDF (full pdf requires viewing direction)
		float exps = WardSmpls[i].x*WardSmpls[i].x/alphax/alphax + WardSmpls[i].y*WardSmpls[i].y/alphay/alphay;
		exps /= hz*hz;

		float norm = 1.f/(4.f*((float) M_PI)*alphax*alphay);

		float pdf = norm/(hz*hz*hz)*expf(-exps); // /dot_hv (done in shader)
		WardSmpls[i].w = pdf;

		WardScales[i] = hz*hz*hz;

	}

}